"""
数据加载和预处理模块

基于reference/1-1-loaddata_preprocess.py改进
实现智能数据筛选和预处理功能
"""

import numpy as np
import pandas as pd
import scipy.io as sio
import scipy.signal as signal
import os
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler, LabelEncoder
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class DataLoader:
    """数据加载器 - 基于reference/1-1改进"""
    
    def __init__(self, source_data_path, target_data_path, target_sampling_rate=12000):
        self.source_data_path = source_data_path
        self.target_data_path = target_data_path
        self.target_fs = target_sampling_rate
        self.scaler = StandardScaler()
        
        # 轴承参数信息
        self.bearing_params = {
            'SKF6205': {'n': 9, 'd': 0.3126, 'D': 1.537},  # 驱动端轴承
            'SKF6203': {'n': 9, 'd': 0.2656, 'D': 1.122}   # 风扇端轴承
        }
        
    def load_source_data(self, sampling_strategy='balanced'):
        """
        加载源域数据并进行智能筛选
        
        Args:
            sampling_strategy: 采样策略 ('balanced', 'representative', 'all')
        """
        print("📂 开始加载源域数据...")
        
        # 获取所有源域文件
        all_files = self._get_all_source_files()
        print(f"发现源域文件总数: {len(all_files)}")
        
        # 分析文件分布
        file_analysis = self._analyze_file_distribution(all_files)
        self._print_file_analysis(file_analysis)
        
        # 根据策略筛选文件
        if sampling_strategy == 'balanced':
            selected_files = self._balanced_sampling(file_analysis)
        elif sampling_strategy == 'representative':
            selected_files = self._representative_sampling(file_analysis)
        elif sampling_strategy == 'systematic':
            selected_files = self._systematic_sampling(file_analysis)
        else:
            selected_files = all_files
            
        print(f"筛选后文件数量: {len(selected_files)}")
        
        # 加载选中的文件
        source_data = self._load_mat_files(selected_files, 'source')
        
        return source_data, selected_files, file_analysis
    
    def load_target_data(self):
        """加载目标域数据"""
        print("📂 开始加载目标域数据...")
        
        target_files = []
        for i in range(16):  # A-P共16个文件
            file_name = f"{chr(65+i)}.mat"  # A.mat, B.mat, ..., P.mat
            file_path = os.path.join(self.target_data_path, file_name)
            if os.path.exists(file_path):
                target_files.append(file_path)
        
        print(f"发现目标域文件数量: {len(target_files)}")
        
        # 加载目标域文件
        target_data = self._load_mat_files(target_files, 'target')
        
        return target_data
    
    def _get_all_source_files(self):
        """获取所有源域文件"""
        all_files = []
        for root, dirs, files in os.walk(self.source_data_path):
            for file in files:
                if file.endswith('.mat'):
                    all_files.append(os.path.join(root, file))
        return all_files
    
    def _analyze_file_distribution(self, files):
        """分析文件分布"""
        analysis = {
            'fault_types': {'B': [], 'IR': [], 'OR': [], 'N': []},
            'fault_sizes': {'0.007': [], '0.014': [], '0.021': [], '0.028': []},
            'load_conditions': {'0': [], '1': [], '2': [], '3': []},
            'sampling_rates': {'12kHz': [], '48kHz': []},
            'positions': {'3点钟': [], '6点钟': [], '12点钟': []}
        }
        
        for file_path in files:
            file_name = os.path.basename(file_path)
            
            # 解析故障类型
            if file_name.startswith('B'):
                analysis['fault_types']['B'].append(file_path)
            elif file_name.startswith('IR'):
                analysis['fault_types']['IR'].append(file_path)
            elif file_name.startswith('OR'):
                analysis['fault_types']['OR'].append(file_path)
            elif file_name.startswith('N'):
                analysis['fault_types']['N'].append(file_path)
            
            # 解析故障尺寸
            for size in ['0.007', '0.014', '0.021', '0.028']:
                if size in file_name:
                    analysis['fault_sizes'][size].append(file_path)
            
            # 解析载荷条件
            for load in ['0', '1', '2', '3']:
                if f'_{load}.mat' in file_name:
                    analysis['load_conditions'][load].append(file_path)
            
            # 解析采样频率
            if '48kHz' in file_name:
                analysis['sampling_rates']['48kHz'].append(file_path)
            else:
                analysis['sampling_rates']['12kHz'].append(file_path)
            
            # 解析外圈故障位置
            if '@3' in file_name:
                analysis['positions']['3点钟'].append(file_path)
            elif '@6' in file_name:
                analysis['positions']['6点钟'].append(file_path)
            elif '@12' in file_name:
                analysis['positions']['12点钟'].append(file_path)
        
        return analysis
    
    def _print_file_analysis(self, analysis):
        """打印文件分析结果"""
        print("\n📊 源域数据分布分析:")
        print("=" * 50)
        
        print("故障类型分布:")
        for fault_type, files in analysis['fault_types'].items():
            print(f"  {fault_type}: {len(files)}个文件")
        
        print("\n故障尺寸分布:")
        for size, files in analysis['fault_sizes'].items():
            print(f"  {size}英寸: {len(files)}个文件")
        
        print("\n载荷条件分布:")
        for load, files in analysis['load_conditions'].items():
            print(f"  {load}马力: {len(files)}个文件")
        
        print("\n采样频率分布:")
        for rate, files in analysis['sampling_rates'].items():
            print(f"  {rate}: {len(files)}个文件")
        
        print("\n外圈故障位置分布:")
        for pos, files in analysis['positions'].items():
            print(f"  {pos}: {len(files)}个文件")
    
    def _balanced_sampling(self, analysis):
        """
        平衡采样策略
        确保每个故障类型都有足够的代表性样本
        """
        print("\n🎯 执行平衡采样策略...")
        
        selected_files = []
        
        # 定义每个故障类型的采样数量
        sampling_plan = {
            'B': 10,    # 滚动体故障：选择10个样本
            'IR': 10,   # 内圈故障：选择10个样本  
            'OR': 15,   # 外圈故障：选择15个样本（包含不同位置）
            'N': 4      # 正常状态：选择4个样本
        }
        
        for fault_type, target_count in sampling_plan.items():
            available_files = analysis['fault_types'][fault_type]
            if len(available_files) >= target_count:
                # 随机选择指定数量的文件
                selected = np.random.choice(available_files, target_count, replace=False)
                selected_files.extend(selected)
                print(f"  {fault_type}: 从{len(available_files)}个文件中选择了{len(selected)}个")
            else:
                # 如果可用文件不足，全部选择
                selected_files.extend(available_files)
                print(f"  {fault_type}: 可用文件不足，选择了全部{len(available_files)}个")
        
        return selected_files
    
    def _systematic_sampling(self, analysis):
        """
        系统性筛选策略 - 基于轴承故障机理
        
        选择原则：
        1. 确保每种故障类型都有代表性样本
        2. 覆盖不同故障尺寸和载荷条件
        3. 考虑外圈故障的不同位置
        4. 优先选择典型工况下的数据
        """
        print("\n🎯 执行系统性筛选策略...")
        print("📋 基于轴承故障机理的科学筛选...")
        
        selected_files = []
        selection_reasoning = {}
        
        # 1. 正常状态数据 - 选择所有可用数据
        normal_files = analysis['fault_types']['N']
        selected_files.extend(normal_files)
        selection_reasoning['Normal'] = {
            'count': len(normal_files),
            'reason': '正常状态样本较少，选择全部数据以确保代表性'
        }
        
        # 2. 滚动体故障 - 选择不同故障尺寸和载荷条件的组合
        ball_files = analysis['fault_types']['B']
        ball_selected = self._select_ball_fault_samples(ball_files)
        selected_files.extend(ball_selected)
        selection_reasoning['Ball'] = {
            'count': len(ball_selected),
            'reason': '选择不同故障尺寸(007,014,021,028)和载荷条件(0,1,2,3)的典型组合'
        }
        
        # 3. 内圈故障 - 选择不同故障尺寸和载荷条件的组合
        inner_files = analysis['fault_types']['IR']
        inner_selected = self._select_inner_fault_samples(inner_files)
        selected_files.extend(inner_selected)
        selection_reasoning['Inner Race'] = {
            'count': len(inner_selected),
            'reason': '选择不同故障尺寸和载荷条件的典型组合，确保内圈故障特征的代表性'
        }
        
        # 4. 外圈故障 - 选择不同位置、故障尺寸和载荷条件的组合
        outer_files = analysis['fault_types']['OR']
        outer_selected = self._select_outer_fault_samples(outer_files)
        selected_files.extend(outer_selected)
        selection_reasoning['Outer Race'] = {
            'count': len(outer_selected),
            'reason': '选择不同位置(3点,6点,12点)、故障尺寸和载荷条件的典型组合'
        }
        
        # 打印选择理由
        self._print_selection_reasoning(selection_reasoning)
        
        return selected_files
    
    def _select_ball_fault_samples(self, ball_files):
        """选择滚动体故障样本"""
        selected = []
        
        # 按故障尺寸和载荷条件分组
        grouped_files = self._group_files_by_conditions(ball_files)
        
        # 选择策略：每个故障尺寸选择2个不同载荷条件的样本
        for size in ['007', '014', '021', '028']:
            if size in grouped_files:
                size_files = grouped_files[size]
                # 按载荷条件分组
                load_groups = {}
                for file_path in size_files:
                    for load in ['0', '1', '2', '3']:
                        if f'_{load}.mat' in file_path:
                            if load not in load_groups:
                                load_groups[load] = []
                            load_groups[load].append(file_path)
                            break
                
                # 选择不同载荷条件的样本
                selected_loads = ['0', '2']  # 选择0马力和2马力作为典型工况
                for load in selected_loads:
                    if load in load_groups and load_groups[load]:
                        selected.append(load_groups[load][0])  # 选择第一个
                        if len(selected) >= 8:  # 限制总数
                            break
                if len(selected) >= 8:
                    break
        
        return selected[:8]  # 最多选择8个
    
    def _select_inner_fault_samples(self, inner_files):
        """选择内圈故障样本"""
        selected = []
        
        # 按故障尺寸和载荷条件分组
        grouped_files = self._group_files_by_conditions(inner_files)
        
        # 选择策略：每个故障尺寸选择2个不同载荷条件的样本
        for size in ['007', '014', '021', '028']:
            if size in grouped_files:
                size_files = grouped_files[size]
                # 按载荷条件分组
                load_groups = {}
                for file_path in size_files:
                    for load in ['0', '1', '2', '3']:
                        if f'_{load}.mat' in file_path:
                            if load not in load_groups:
                                load_groups[load] = []
                            load_groups[load].append(file_path)
                            break
                
                # 选择不同载荷条件的样本
                selected_loads = ['0', '2']  # 选择0马力和2马力作为典型工况
                for load in selected_loads:
                    if load in load_groups and load_groups[load]:
                        selected.append(load_groups[load][0])  # 选择第一个
                        if len(selected) >= 8:  # 限制总数
                            break
                if len(selected) >= 8:
                    break
        
        return selected[:8]  # 最多选择8个
    
    def _select_outer_fault_samples(self, outer_files):
        """选择外圈故障样本"""
        selected = []
        
        # 按位置分组
        position_groups = {
            '@3': [f for f in outer_files if '@3' in f],
            '@6': [f for f in outer_files if '@6' in f],
            '@12': [f for f in outer_files if '@12' in f]
        }
        
        # 从每个位置选择代表性样本
        for position, files in position_groups.items():
            if not files:
                continue
                
            # 按故障尺寸分组
            grouped_files = self._group_files_by_conditions(files)
            
            # 选择不同故障尺寸的样本
            for size in ['007', '014', '021']:
                if size in grouped_files and grouped_files[size]:
                    # 选择0马力载荷条件的样本
                    for file_path in grouped_files[size]:
                        if '_0.mat' in file_path:
                            selected.append(file_path)
                            break
                    if len(selected) >= 12:  # 限制总数
                        break
                if len(selected) >= 12:
                    break
            if len(selected) >= 12:
                break
        
        return selected[:12]  # 最多选择12个
    
    def _group_files_by_conditions(self, files):
        """按故障尺寸分组文件"""
        grouped = {}
        for file_path in files:
            file_name = file_path.split('\\')[-1]  # 获取文件名
            for size in ['007', '014', '021', '028']:
                if size in file_name:
                    if size not in grouped:
                        grouped[size] = []
                    grouped[size].append(file_path)
                    break
        return grouped
    
    def _print_selection_reasoning(self, reasoning):
        """打印选择理由"""
        print("\n📋 数据筛选理由说明:")
        print("=" * 50)
        
        for fault_type, info in reasoning.items():
            print(f"\n{fault_type}故障:")
            print(f"  选择数量: {info['count']}个样本")
            print(f"  选择理由: {info['reason']}")
        
        print("\n🎯 总体筛选原则:")
        print("1. 基于轴承故障机理选择典型工况")
        print("2. 确保每种故障类型都有代表性样本")
        print("3. 覆盖不同故障尺寸和载荷条件")
        print("4. 考虑外圈故障的不同位置特征")
        print("5. 优先选择0马力和2马力等典型工况")
    
    def _representative_sampling(self, analysis):
        """
        代表性采样策略
        选择最能代表各类故障特征的文件
        """
        print("\n🎯 执行代表性采样策略...")
        
        selected_files = []
        
        # 为每个故障类型选择代表性样本
        for fault_type in ['B', 'IR', 'OR', 'N']:
            available_files = analysis['fault_types'][fault_type]
            if not available_files:
                continue
                
            # 选择不同故障尺寸和载荷条件的组合
            representative_files = self._select_representative_files(available_files, fault_type)
            selected_files.extend(representative_files)
            print(f"  {fault_type}: 选择了{len(representative_files)}个代表性文件")
        
        return selected_files
    
    def _select_representative_files(self, files, fault_type):
        """为特定故障类型选择代表性文件"""
        if fault_type == 'OR':
            # 外圈故障：选择不同位置的代表性文件
            selected = []
            for pos in ['@3', '@6', '@12']:
                pos_files = [f for f in files if pos in f]
                if pos_files:
                    # 选择不同故障尺寸的文件
                    for size in ['007', '014', '021']:
                        size_files = [f for f in pos_files if size in f]
                        if size_files:
                            selected.append(size_files[0])  # 选择第一个
            return selected[:5]  # 最多选择5个
        
        elif fault_type in ['B', 'IR']:
            # 滚动体和内圈故障：选择不同故障尺寸和载荷条件的组合
            selected = []
            for size in ['007', '014', '021', '028']:
                size_files = [f for f in files if size in f]
                if size_files:
                    # 选择不同载荷条件的文件
                    for load in ['0', '1', '2', '3']:
                        load_files = [f for f in size_files if f'_{load}.mat' in f]
                        if load_files:
                            selected.append(load_files[0])  # 选择第一个
                            break  # 每个尺寸只选择一个
            return selected[:5]  # 最多选择5个
        
        else:  # 正常状态
            return files[:4]  # 选择前4个
    
    def _load_mat_files(self, file_list, data_type):
        """加载MAT文件数据"""
        all_data = []
        
        for file_path in tqdm(file_list, desc=f"Loading {data_type} files"):
            try:
                data = sio.loadmat(file_path)
                file_name = os.path.basename(file_path)
                
                # 提取信号数据
                de_signal = self._extract_signal(data, 'DE')
                fe_signal = self._extract_signal(data, 'FE')
                ba_signal = self._extract_signal(data, 'BA')
                rpm = self._extract_rpm(data)
                
                # 提取文件信息
                file_info = self._parse_file_info(file_name, data_type)
                
                all_data.append({
                    'file_name': file_name,
                    'de_signal': de_signal,
                    'fe_signal': fe_signal,
                    'ba_signal': ba_signal,
                    'rpm': rpm,
                    'sampling_rate': file_info['sampling_rate'],
                    'fault_type': file_info['fault_type'],
                    'fault_size': file_info['fault_size'],
                    'load_condition': file_info['load_condition'],
                    'data_type': data_type
                })
                
            except Exception as e:
                print(f"⚠️ 加载文件失败 {file_path}: {e}")
                continue
        
        print(f"✅ 成功加载 {len(all_data)} 个{data_type}文件")
        return all_data
    
    def _extract_signal(self, data, signal_type):
        """提取信号数据"""
        # 查找包含信号类型的键
        for key in data.keys():
            if signal_type in key and 'time' in key:
                signal_data = data[key].flatten()
                # 重采样到目标采样率
                if len(signal_data) > 0:
                    return self._resample_signal(signal_data, signal_type)
        
        # 对于目标域数据，尝试直接使用单字母键
        if signal_type == 'DE':
            for key in data.keys():
                if len(key) == 1 and key.isalpha():  # 单字母键如'A', 'B'等
                    signal_data = data[key].flatten()
                    if len(signal_data) > 0:
                        return self._resample_signal(signal_data, signal_type)
        
        # 如果没找到，返回空数组
        return np.array([])
    
    def _resample_signal(self, signal, signal_type):
        """重采样信号到目标采样率"""
        if len(signal) == 0:
            return signal
            
        # 根据信号类型确定原始采样率
        if signal_type == 'DE':
            original_fs = 12000  # 驱动端12kHz
        else:
            original_fs = 12000  # 风扇端和基座12kHz
        
        if original_fs == self.target_fs:
            return signal
        
        # 计算重采样参数
        num_samples = int(len(signal) * self.target_fs / original_fs)
        resampled_signal = signal.resample(signal, num_samples)
        
        return resampled_signal
    
    def _extract_rpm(self, data):
        """提取RPM数据"""
        for key in data.keys():
            if 'rpm' in key.lower():
                return float(data[key][0][0])
        return 1796.0  # 默认RPM
    
    def _parse_file_info(self, file_name, data_type):
        """解析文件信息"""
        if data_type == 'target':
            return {
                'sampling_rate': 32000,  # 目标域32kHz
                'fault_type': 'Unknown',
                'fault_size': 0.0,
                'load_condition': 0
            }
        
        # 解析源域文件信息
        sampling_rate = 12000
        if '48kHz' in file_name:
            sampling_rate = 48000
        
        fault_type = 'Normal'
        if file_name.startswith('B'):
            fault_type = 'Ball'
        elif file_name.startswith('IR'):
            fault_type = 'Inner Race'
        elif file_name.startswith('OR'):
            fault_type = 'Outer Race'
        
        fault_size = 0.0
        for size in ['007', '014', '021', '028']:
            if size in file_name:
                fault_size = float(size) / 1000
                break
        
        load_condition = 0
        for load in ['0', '1', '2', '3']:
            if f'_{load}.mat' in file_name:
                load_condition = int(load)
                break
        
        return {
            'sampling_rate': sampling_rate,
            'fault_type': fault_type,
            'fault_size': fault_size,
            'load_condition': load_condition
        }
    
    def save_data_separately(self, source_data, target_data, output_dir='.'):
        """分别保存源域和目标域数据"""
        print("\n💾 保存数据文件...")
        
        # 保存源域数据
        source_df = pd.DataFrame(source_data)
        source_file = os.path.join(output_dir, 'source_domain_data.csv')
        source_df.to_csv(source_file, index=False, encoding='utf-8-sig')
        print(f"✅ 源域数据已保存: {source_file}")
        
        # 保存目标域数据
        target_df = pd.DataFrame(target_data)
        target_file = os.path.join(output_dir, 'target_domain_data.csv')
        target_df.to_csv(target_file, index=False, encoding='utf-8-sig')
        print(f"✅ 目标域数据已保存: {target_file}")
        
        return source_file, target_file
